#===========================================================================================================#
# 使用SimNIBS进行tDCS优化
# v0.0.2 2024-11-05
# debug, 
# now default output_path is m2m_{subject_name}/optimize/optimize_resultsYYYYMMDDHHMMSS
# v0.0.5 2024-12-18
# modify the output_path, now will be saved in the output_path/optimize_resultsYYYYMMDDHHMMSS.msh instead of output_path.msh,
# add post flag(whether to post into .nii file, default is True)
# example:
# --ldfd_hdf = "C:\Users\liushiyu0523\Downloads\simnibs4_examples\leadfield\ernie_leadfield_EEG10-10_UI_Jurak_2007.hdf5"
# --output_path = "C:\Users\liushiyu0523\Downloads\simnibs4_examples\m2m_ernie\output_20241218"
# --target_image = "C:\Users\liushiyu0523\Downloads\simnibs4_examples\m2m_ernie\T1.nii.gz"
# --subject_folder = "C:\Users\liushiyu0523\Downloads\simnibs4_examples\m2m_ernie"
# --leadfield_folder_name = "C:\Users\liushiyu0523\Downloads\simnibs4_examples\leadfield"
# --silent 
# --post False
# cd C:\Users\liushiyu0523\guangdianganying_software # where python files are 
# python tdcs_optimize_el_network.py --ldfd_hdf "C:\Users\liushiyu0523\Downloads\simnibs4_examples\leadfield\ernie_leadfield_EEG10-10_UI_Jurak_2007.hdf5" --output_path "C:\Users\liushiyu0523\Downloads\simnibs4_examples\m2m_ernie\output_20241218" --target_image "C:\Users\liushiyu0523\Downloads\simnibs4_examples\m2m_ernie\T1.nii.gz" --subject_folder "C:\Users\liushiyu0523\Downloads\simnibs4_examples\m2m_ernie" --leadfield_folder_name "C:\Users\liushiyu0523\Downloads\simnibs4_examples\leadfield" --silent --no-post
# ### output
# =====================================
# SimNIBS finished running optimization
# Mesh file: C:\Users\liushiyu0523\Downloads\simnibs4_examples\m2m_ernie\output_20241218\optimize_results20241218121906.msh
# CSV file: C:\Users\liushiyu0523\Downloads\simnibs4_examples\m2m_ernie\output_20241218\optimize_results20241218121906.csv
# Summary file: C:\Users\liushiyu0523\Downloads\simnibs4_examples\m2m_ernie\output_20241218\optimize_results20241218121906_summary.txt
# =====================================
# no_post: True
# No post-process
#===========================================================================================================#

import simnibs
import os
import subprocess
import platform
from simnibs.utils import transformations
import multiprocessing
import argparse
from datetime import datetime

# Post-processing functions
# def post_process_silent(opt_name, subject_folder):
#     transformations.middle_gm_interpolation(
#         opt_name + ".msh", subject_folder, opt_name + '_subject_overlays',
#         out_fsaverage=opt_name + '_fsavg_overlays', depth=0.5,
#         open_in_gmsh=False, f_geo=opt_name + "_el_currents.geo"
#     )
#     transformations.interpolate_to_volume(
#         opt_name + ".msh", subject_folder, opt_name, keep_tissues=[1, 2, 3, 6]
#     )
#     transformations.warp_volume(
#         opt_name + ".msh", subject_folder, opt_name, keep_tissues=[1, 2, 3, 6]
#     )

# def post_process(opt_name, subject_folder):
#     transformations.middle_gm_interpolation(
#         opt_name + ".msh", subject_folder, opt_name + '_subject_overlays',
#         out_fsaverage=opt_name + '_fsavg_overlays', depth=0.5,
#         open_in_gmsh=True, f_geo=opt_name + "_el_currents.geo"
#     )
#     transformations.interpolate_to_volume(
#         opt_name + ".msh", subject_folder, opt_name, keep_tissues=[1, 2, 3, 6]
#     )
#     transformations.warp_volume(
#         opt_name + ".msh", subject_folder, opt_name, keep_tissues=[1, 2, 3, 6]
#     )
#     # Open output in ITK-SNAP
#     module = "itk-snap" if platform.system() == 'Windows' else "itksnap"
#     subprocess.run([module, '-g', f'{opt_name}_magnE.nii.gz'])
#     subprocess.run([module, '-g', f'{opt_name}_MNI_magnE.nii.gz'])
from tdcs_optimize_el import post_process_silent,post_process

# Simulation function
def run_simulation(subject_folder, target_image,output_path=None, ldfd_hdf=None,target_image_in_mni=False,
                   target_intensity=None, silent=False,
                   max_total_current=2e-3,max_individual_current=1e-3,max_active_electrodes=8,
                   leadfield_folder_name="leadfield",no_post=False):
    if ldfd_hdf==None:
        ldfd_file=os.path.basename(subject_folder)[4:]+"_leadfield_EEG10-10_UI_Jurak_2007.hdf5"
        ldfd_hdf=os.path.join(subject_folder,leadfield_folder_name,ldfd_file)
        print("Using default path:\n",ldfd_hdf)
    if not os.path.exists(ldfd_hdf):
        print("leadfield_hdf not exists, check if exits or wrong path:\n",ldfd_hdf)
    if output_path==None:
        optimize_dir=os.path.join(subject_folder,"optimize")
        if not os.path.exists(optimize_dir):
            os.mkdir(optimize_dir)
        output_path=os.path.join(optimize_dir,"optimize_results"+datetime.now().strftime('%Y%m%d%H%M%S'))
        print("Using default output_path:\n",output_path)
    else:
        output_path=os.path.join(output_path,"optimize_results"+datetime.now().strftime('%Y%m%d%H%M%S'))
    if target_intensity==None:
        target_intensity=0.2
        print("Using default intensity for target:\n",target_intensity)
    
    # opt = simnibs.opt_struct.TDCSoptimize()
    # opt.leadfield_hdf = ldfd_hdf
    # opt.name = output_path
    # opt.max_total_current = max_total_current
    # opt.max_individual_current = max_individual_current
    # opt.max_active_electrodes = max_active_electrodes
    
    opt = simnibs.opt_struct.TDCSDistributedOptimize()
    opt.leadfield_hdf = ldfd_hdf
    opt.subpath = subject_folder
    opt.name = output_path
    opt.max_total_current = max_total_current
    opt.max_individual_current = max_individual_current
    opt.max_active_electrodes = max_active_electrodes
    opt.target_image = target_image
    opt.mni_space = target_image_in_mni
    opt.intensity = target_intensity
    print("target_image",target_image)
    simnibs.run_simnibs(opt)
    
    # Run post-process functions in parallel
    print("no_post:",no_post)
    if not no_post:
        opt_name = output_path
        process_func = post_process_silent if silent else post_process
        p = multiprocessing.Process(target=process_func, args=(opt_name, subject_folder))
        p.start()
    else:
        print("No post-process")

# Main function with argument parsing
def main():
    parser = argparse.ArgumentParser(description="Run tDCS optimization using SimNIBS")
    parser.add_argument('--ldfd_hdf', help="Path to leadfield HDF file")
    parser.add_argument('--output_path', help="Output path for simulation results")
    #parser.add_argument('--target_image', required=True, type=float, nargs=3, help="Target coordinates,for example,ID03_MOTOR_ICA.nii.gz")
    parser.add_argument('--target_image', required=True, help="Target image path")
    parser.add_argument('--subject_folder', required=True, help="Subject folder path")
    parser.add_argument('--target_image_in_mni', action='store_true', help="Set if target images are in MNI space")
    parser.add_argument('--target_intensity', type=float, default=0.2, help="Target intensity (default: 0.2)")
    parser.add_argument('--silent', action='store_true', help="Run in silent mode (no visualization)")
    parser.add_argument('--max_total_current', help="Indicates max_total_current")
    parser.add_argument('--max_individual_current', help="Indicates max_individual_current")
    parser.add_argument('--max_active_electrodes',  help="Indicates max_active_electrodes'")
    # add leadfield_folder_name,default is leadfield
    parser.add_argument('--leadfield_folder_name', default="leadfield",required=True, help="Leadfield folder name")
    # add post, default is True
    parser.add_argument('--no-post', action='store_true', help="Not run post-process,default is run")
    args = parser.parse_args()
    
    run_simulation(
        ldfd_hdf=args.ldfd_hdf,
        output_path=args.output_path,
        #target_coord=args.target_coord,
        target_image=args.target_image,
        subject_folder=args.subject_folder,
        #target_coord_in_mni=args.target_coord_in_mni,
        target_image_in_mni=args.target_image_in_mni,
        target_intensity=args.target_intensity,
        #avoid_coord=args.avoid_coord,
        #avoid_coord_in_mni=args.avoid_coord_in_mni,
        silent=args.silent,
        leadfield_folder_name=args.leadfield_folder_name,
        #ldfd_hdf_volume=args.ldfd_hdf_volume
        no_post=args.no_post
    )

if __name__ == "__main__":
    main()
